#!/usr/bin/env python
import rospy
import cv2
import numpy as np
from sensor_msgs.msg import Image
from std_msgs.msg import String
from cv_bridge import CvBridge

class FireDetector:
    def __init__(self):
        rospy.init_node("fire_detection_node", anonymous=True)

        # 订阅无人机摄像头图像
        self.image_sub = rospy.Subscriber("/usb_cam/image_raw", Image, self.image_callback)

        # 发布火源检测信息 & 可视化图像
        self.info_pub = rospy.Publisher("/AI_detect_info", String, queue_size=1)
        self.image_pub = rospy.Publisher("/fire_detect_image", Image, queue_size=1)

        self.bridge = CvBridge()

        # **更新 HSV 阈值（你提供的新参数）**
        self.lower_hsv = np.array([0, 143, 62])    # H_min, S_min, V_min
        self.upper_hsv = np.array([179, 255, 255])  # H_max, S_max, V_max

    def image_callback(self, msg):
        try:
            frame = self.bridge.imgmsg_to_cv2(msg, "bgr8")
        except Exception as e:
            rospy.logerr("cv_bridge 错误: %s", str(e))
            # 若未找到目标，发送 -1 -1 坐标
            detection_info = " not_found_fire!!!: x:{} y:{};".format(-1, -1)
            self.info_pub.publish(detection_info)
            return

        if frame is None:
            rospy.logwarn("接收到空图像帧")
            # 若未找到目标，发送 -1 -1 坐标
            detection_info = " not_found_fire!!!: x:{} y:{};".format(-1, -1)
            self.info_pub.publish(detection_info)
            return

        # **检查分辨率**
        h, w, _ = frame.shape
        if w != 1280 or h != 720:
            rospy.logwarn("图像分辨率错误: 期望 1280x720, 但收到 %dx%d", w, h)
            # 若未找到目标，发送 -1 -1 坐标
            detection_info = " not_found_fire!!!: x:{} y:{};".format(-1, -1)
            self.info_pub.publish(detection_info)
            return

        # **电子裁剪 (中央 640x480)**
        crop_x = (1280 - 640) // 2  # 320
        crop_y = (720 - 480) // 2   # 120
        cropped_frame = frame[crop_y:crop_y + 480, crop_x:crop_x + 640]

        # **转换到 HSV 颜色空间**
        hsv = cv2.cvtColor(cropped_frame, cv2.COLOR_BGR2HSV)

        # **颜色筛选 (使用你的 HSV 阈值)**
        mask = cv2.inRange(hsv, self.lower_hsv, self.upper_hsv)

        # **形态学处理 (去除噪点)**
        mask = cv2.erode(mask, None, iterations=2)
        mask = cv2.dilate(mask, None, iterations=2)

        # **轮廓检测（兼容 OpenCV 3.x 和 4.x）**
        contour_results = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        if len(contour_results) == 2:
            contours, _ = contour_results  # OpenCV 3.x
        else:
            _, contours, _ = contour_results  # OpenCV 4.x

        # 初始化为未找到目标的信息
        detection_info = " not_found_fire!!!: x:{} y:{};".format(-1, -1)

        # **只处理最大轮廓**
        if contours:
            max_contour = max(contours, key=cv2.contourArea)
            area = cv2.contourArea(max_contour)

            if area >= 500:  # 过滤掉小目标
                x, y, w, h = cv2.boundingRect(max_contour)
                center_x = x + w // 2
                center_y = y + h // 2

                # **绘制矩形和中心点**
                cv2.rectangle(cropped_frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
                cv2.circle(cropped_frame, (center_x, center_y), 5, (0, 255, 0), -1)

                # **按 Yolo 格式发布火源信息**
                detection_info = "found_fire!!!: x:{} y:{};".format(center_x, center_y)
                print("found_fire!!!")
                rospy.loginfo(detection_info)

        # **发布检测结果**
        self.info_pub.publish(detection_info)

        # **发布可视化图像**
        vis_msg = self.bridge.cv2_to_imgmsg(cropped_frame, "bgr8")
        self.image_pub.publish(vis_msg)

if __name__ == '__main__':
    try:
        FireDetector()
        rospy.spin()
    except rospy.ROSInterruptException:
        pass
